/**
 *
 * DO NOT EDIT THIS FILE UNLESS IT IS IN /costs/src/index.ts
 */

// TODO this can all move to Valhalla in Arch-ReDesign V2

import { costs as openaiCosts } from "./providers/openai";
import { costs as fineTunedOpenAICosts } from "./providers/openai/fine-tuned-models";
import { costs as togetherAICosts } from "./providers/togetherai";
import { costs as azureCosts } from "./providers/azure";
import { costs as googleCosts } from "./providers/google";

export function costOf(model: string) {
  const costs = [
    ...openaiCosts,
    ...togetherAICosts,
    ...azureCosts,
    ...googleCosts,
    ...fineTunedOpenAICosts,
  ];

  const cost = costs.find((cost) => {
    if (cost.model.operator === "equals") {
      return cost.model.value === model;
    } else if (cost.model.operator === "startsWith") {
      return model.startsWith(cost.model.value);
    }
  });

  return cost?.cost;
}

export function costOfPrompt({
  model,
  promptTokens,
  completionTokens,
}: {
  model: string;
  promptTokens: number;
  completionTokens: number;
}) {
  const cost = costOf(model);
  if (!cost) {
    return 0;
  }
  return (
    cost.prompt_token * promptTokens + cost.completion_token * completionTokens
  );
}

export function clickhousePriceCalc(table: string) {
  const costs = [
    ...openaiCosts,
    ...togetherAICosts,
    ...azureCosts,
    ...googleCosts,
    ...fineTunedOpenAICosts,
  ];
  // This is so that we don't need to do any floating point math in the database
  // and we can just divide by 1_000_000 to get the cost in dollars
  const multiple = 1_000_000_000;
  return `
sum(
  CASE
    ${costs
      .map((cost) => {
        if (cost.model.operator === "equals") {
          return `WHEN (${table}.model = '${cost.model.value}') THEN ${
            cost.cost.prompt_token * multiple
          } * ${table}.prompt_tokens + ${
            cost.cost.completion_token * multiple
          } * ${table}.completion_tokens`;
        } else if (cost.model.operator === "startsWith") {
          return `WHEN (${table}.model LIKE '${cost.model.value}%') THEN ${
            cost.cost.prompt_token * multiple
          } * ${table}.prompt_tokens + ${
            cost.cost.completion_token * multiple
          } * ${table}.completion_tokens`;
        }
      })
      .join("\n")}
    ELSE 0
  END
  ) / ${multiple}
`;
}
